from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QMenu, QTreeView, QDirModel, QAction
from PyQt5.QtCore import Qt, QDir, QFileInfo, QModelIndex
import qtawesome as qta
import os


class FileTreeItem(QTreeWidgetItem):
    def __init__(self):
        super(FileTreeItem, self).__init__()
        self.file_path = ''
        self.item_type = 'dir'   # computer---driver---dir---file


class FileTreeWidget(QTreeWidget):
    """
    一个类似于文件浏览器的树形部件，
    """

    def __init__(self):
        super(FileTreeWidget, self).__init__()

        self.setHeaderHidden(True)

        self.root = FileTreeItem()
        self.setRootItem()

        self.contextMenu = QMenu()
        # self.contextActionDict = {'computer': [], 'driver': [], 'dir': []}
        self.contextMenuDict = {'computer': QMenu(), 'driver': QMenu(), 'dir': QMenu()}
        self.openAction = QAction('展开')
        self.deleteAction = QAction('删除')
        self.renameAction = QAction('重命名')
        self.propertyAction = QAction('属性')

        self.createContextMenu()
        # self.tree_model = QDirModel()
        # self.setModel(self.tree_model)
        # self.itemCollapsed.connect()

    def setRootItem(self):
        # 显示所有文件夹
        self.root.setIcon(0, qta.icon('mdi.desktop-classic'))
        self.root.setText(0, '此电脑')
        self.root.item_type = 'computer'
        self.root.file_path = '/'

        self.addTopLevelItem(self.root)

        # for root, dirs, files in os.walk('/', topdown=True):
            # for name in files:
            #    print(os.path.join(root, name))
        #    for name in dirs:
        #        print(os.path.join(root, name))

        drivers = QDir().drives()
        # print(drivers)
        for item in drivers:
            tree_item = FileTreeItem()
            tree_item.item_type = 'driver'
            tree_item.file_path = item.absolutePath()
            tree_item.setText(0, item.absolutePath())
            self.root.addChild(tree_item)
            self.addChildren(tree_item)

    def addChildren(self, parent_item: FileTreeItem):

        info_List = QDir(parent_item.file_path).entryInfoList()

        for item in info_List:
            tree_item = FileTreeItem()
            if item.isDir() and item.fileName() != '.' and item.fileName() != '..':
                tree_item.setIcon(0, qta.icon('mdi.folder', color='dodgerblue'))
                tree_item.setText(0, item.fileName())
                tree_item.file_path = item.filePath()
                tree_item.item_type = 'dir'
                parent_item.addChild(tree_item)
                # if item.
                # self.addChildren(tree_item)

    def createContextMenu(self):

        self.openAction.triggered.connect(self.open_folder)
        self.deleteAction.triggered.connect(self.deleteLater)
        self.propertyAction.triggered.connect(self.property)

        # self.contextActionDict['computer'].append(self.openAction)
        # self.contextActionDict['computer'].append(self.propertyAction)
        #
        # self.contextActionDict['driver'].append(self.openAction)
        # self.contextActionDict['driver'].append(self.propertyAction)
        #
        # self.contextActionDict['dir'].append(self.openAction)
        # self.contextActionDict['dir'].append(self.deleteAction)
        # self.contextActionDict['dir'].append(self.renameAction)
        # self.contextActionDict['dir'].append(self.propertyAction)

        self.contextMenuDict['computer'].addAction(self.openAction)
        self.contextMenuDict['computer'].addAction(self.propertyAction)
        self.contextMenuDict['driver'].addAction(self.openAction)
        self.contextMenuDict['driver'].addAction(self.propertyAction)

        self.contextMenuDict['dir'].addAction(self.openAction)
        self.contextMenuDict['dir'].addSeparator()
        self.contextMenuDict['dir'].addAction(self.deleteAction)
        self.contextMenuDict['dir'].addAction(self.renameAction)
        self.contextMenuDict['dir'].addSeparator()
        self.contextMenuDict['dir'].addAction(self.propertyAction)

        self.setContextMenuPolicy(Qt.CustomContextMenu)
        self.customContextMenuRequested.connect(self.showContextMenu)

    def showContextMenu(self, pos):
        if self.currentItem().isExpanded():
            self.openAction.setText('折叠')
        else:
            self.openAction.setText('展开')
        item_type = self.currentItem().item_type
        menu = self.contextMenuDict[item_type]
        menu.move(self.mapToGlobal(pos))
        menu.show()

    def open_folder(self):
        if not self.currentItem().isExpanded():
            if self.currentItem().childCount() == 0:

                path = self.currentItem().file_path
                item = QFileInfo()

                info_list = QDir(path).entryInfoList()
                info_list.sort(key=lambda x: x.isDir(), reverse=True)
                for item in info_list:
                    tree_item = FileTreeItem()
                    if item.isDir() and item.fileName() != '.' and item.fileName() != '..':
                        tree_item.setIcon(0, qta.icon('mdi.folder', color='dodgerblue'))
                        tree_item.setText(0, item.fileName())
                        tree_item.file_path = item.filePath()
                        tree_item.item_type = 'dir'
                        self.currentItem().addChild(tree_item)

            self.currentItem().setExpanded(True)
            self.openAction.setText('折叠')
        else:
            self.currentItem().setExpanded(False)
            self.openAction.setText('展开')
        # print(os.listdir(data))



